#include <iostream>
#include <fstream>
#include <string>
#include <unordered_set>
#include <utility>
#include <stdio.h>
#include <stdlib.h>
#include <queue>
#include <chrono>
#include "hnswlib.hpp"
#include "vecs_io.hpp"
#include "space_ip.hpp"
#include "hnswalg.hpp"
#include <algorithm>
#include <ctime>
#include <cassert>
#include <omp.h>

using namespace std;
using namespace MIPSGraph;

const int kTopK = 10;

std::vector<std::priority_queue<std::pair<float, labeltype >>> loadLSHBOX(std::string inputPath) {
    std::vector<std::priority_queue<std::pair<float, labeltype >>> answers;

    std::ifstream fin(inputPath.c_str());
    unsigned numQueries;
    unsigned K;
    fin >> numQueries >> K;
    answers.resize(numQueries);

    unsigned qId;
    unsigned id;
    float dist;
    int index = 0;
    for (int q = 0; q < numQueries; ++q) {
        fin >> qId;
        assert(qId == q);
        for (int i = 0; i < K; ++i) {
            fin >> id >> dist;
            answers[q].emplace(dist, id);
        }
    }
    fin.close();
    return answers;
}

struct ResultPair {
    int efSearch_;
    double time_;
    double recall_;
    double dist_calc_;
};

int main(int argc, char **argv) {
    if (argc != 5) {
        std::printf("query_fname, gnd_fname, inputGraph_fname, output_fname\n");
        return 0;
    }

    std::string queryname = std::string(argv[1]);
    std::string benchmarkname = std::string(argv[2]);
    std::string graphname = std::string(argv[3]);
    std::string outputname = std::string(argv[4]);

    std::vector<std::priority_queue<std::pair<float, labeltype >>> answers = loadLSHBOX(benchmarkname);

    unsigned int qsize, vecdim;
    float *massQ = MIPSGraph::loadVecs<float>(queryname.c_str(), qsize, vecdim);
    MIPSGraph::InnerProductSpace innerProductSpace(vecdim);

    std::priority_queue<std::pair<float, labeltype >> estimate_l[qsize];

    MIPSGraph::HierarchicalNSW<float> *appr_alg = new MIPSGraph::HierarchicalNSW<float>(&innerProductSpace,
                                                                                        graphname.data(),
                                                                                        false);
// =========== very important here ========================================================================================================
    linklistsizeint *temp_data1 = NULL;
    int degree_count = 0;
    float norm_avg = 0, norm_var = 0;
    for (int i = 0; i < appr_alg->maxelements_; ++i) {
        temp_data1 = (linklistsizeint * )(appr_alg->data_level0_memory_ + i * appr_alg->size_data_per_element_);
        degree_count += *temp_data1;
        norm_avg += appr_alg->elementNorms[i];
    }
    norm_avg /= appr_alg->maxelements_;
    for (int i = 0; i < appr_alg->maxelements_; ++i) {
        norm_var += (appr_alg->elementNorms[i] - norm_avg) * (appr_alg->elementNorms[i] - norm_avg);
    }
    norm_var /= appr_alg->maxelements_;
    std::cout << "avg. level0 degree = " << (float) degree_count / appr_alg->maxelements_ << std::endl;
    std::cout << "norm avg. = " << norm_avg << std::endl;
    std::cout << "norm std. = " << sqrt(norm_var) << std::endl;
    // ========================================================================================================================================

    std::cout << "max level : " << appr_alg->maxlevel_ << std::endl;
    std::cout << "ef, " << "recall, " << " query_time, " << "dist_comp" << std::endl;
    std::vector<int> efs;
    for (int i = 10; i < 100; i += 10) {
        efs.push_back(i);
    }
    for (int i = 100; i < 300; i += 20) {
        efs.push_back(i);
    }
    for (int i = 300; i < 1000; i += 100) {
        efs.push_back(i);
    }
//    for (int i = 2000; i < 20000; i += 2000) {
//        efs.push_back(i);
//    }
//    for (int i = 20000; i < 100000; i += 20000) {
//        efs.push_back(i);
//    }

    std::vector<ResultPair> time_recall_l;

    for (int efSearch : efs) {
        appr_alg->setEf(efSearch);
        appr_alg->dist_calc = 0;

        auto start = std::chrono::high_resolution_clock::now();
        for (int i = 0; i < qsize; i++) {
            estimate_l[i] = appr_alg->searchKnnNoHier(massQ + vecdim * i, kTopK);
        }
        auto end = std::chrono::high_resolution_clock::now();
        // std::cout << "quality : " << appr_alg->quality_of_first_bot_bucket / qsize << std::endl;

        int correct = 0, total = 0;

        for (int i = 0; i < qsize; i++) {
            std::vector<int> res;
            while (!estimate_l[i].empty()) {
                res.push_back(estimate_l[i].top().second);
                estimate_l[i].pop();
            }
            std::reverse(res.begin(), res.end());

            std::unordered_set<labeltype> g;
            std::priority_queue<std::pair<float, labeltype >> gnd(answers[i]);
            total += gnd.size();
            while (gnd.size()) {
                g.insert(gnd.top().second);
                gnd.pop();
            }
            for (auto it : res) {
                // to be done soon
                if (g.find(it) != g.end()) {
                    correct++;
                }
            }
        }
        double recall = 1.0f * correct / total;
        double time = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() / (double) qsize;
        double dist_calc = appr_alg->dist_calc / (double) qsize;
        std::cout << efSearch << ", " << recall << ", " << time << ", " << dist_calc << std::endl;
        ResultPair resultPair;
        resultPair.efSearch_ = efSearch;
        resultPair.time_ = time;
        resultPair.recall_ = recall;
        resultPair.dist_calc_ = dist_calc;
        time_recall_l.push_back(resultPair);
    }

    std::ofstream fres;
    fres.open(outputname);
    fres << "ef, recall, time, dist_calc" << std::endl;
    for (int i = 0; i < (int) efs.size(); i++) {
        fres << time_recall_l[i].efSearch_ << ", ";
        fres << time_recall_l[i].recall_ << ", ";
        fres << time_recall_l[i].time_ << ", ";
        fres << time_recall_l[i].dist_calc_ << std::endl;
    }
    fres.close();
    delete appr_alg;
    delete massQ;
    return 0;
}
